import torch
from torch.utils.data import Dataset, DataLoader
import os
import json
from pathlib import Path
import cv2
import numpy as np
import argparse
from hub.yolov8 import yolov8, yolov8_change_head, yolov8_tfhead, yolov8_transfromer_head
from torchvision.ops import nms
from util.common import image_preprocess
from util.metrics import NMS, min_NMS
from dataset.make_dataset import cut_image
from tqdm import tqdm
from torchmetrics.detection.mean_ap import MeanAveragePrecision

def get_submit(images_path, model_weight, export_resault):
    classes = ['car']
    input_shape = [640, 640]
    num_classes = len(classes)
    print(num_classes)
    device = torch.device('cuda')
    net = yolov8.yolov8_detect(nc=num_classes, phi='n').to(device)
    # net = yolov8_change_head.yolov8_detect(nc=num_classes, model_scale=0.33).to(device)
    # net = yolov8_tfhead.yolov8_detect(nc=num_classes, model_scale=0.33).to(device)
    
    net.load_state_dict(torch.load(model_weight, map_location=device))
    # net = net.fuse().eval()    
    net.eval()

    images = os.listdir(images_path)
    submit_dicts = []
    pbar = tqdm(images)
    for image_id in pbar:
        pbar.set_description('get json | {}'.format(image_id))
        with torch.no_grad():
            image_path = os.path.join(images_path, image_id)
            src = cv2.imread(image_path)
            img = src.copy()
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            # src = image_preprocess(src, (800, 800), False)
            img = image_preprocess(img, (640, 640), True)

            img_tensor = torch.from_numpy(img).float().to(device)
            img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0)

            y = net(img_tensor).detach().cpu().squeeze(0)
            y = y.permute(1, 0)
            bbox, cls_pref = y.split((4, num_classes), 1)
            cls_conf, cls_pre = torch.max(cls_pref, 1, keepdim=True)
            dbox = torch.concat((bbox, cls_conf, cls_pre), dim=1)
            bbox, scores, _ = NMS(dbox, scores_threshold=0.5, device=torch.device('cpu'))
            car_box = []
            scale = 800.0 / 640.0
            for i, box in enumerate(bbox):
                # src = cv2.rectangle(src, (int(box[0] * scale), int(box[1] * scale)), (int(box[2] * scale), int(box[3] * scale)), (0, 255, 0, 0.1))
                # src = cv2.putText(src, '{:1.3f}'.format(float(scores[i])),(int(box[0] * scale), int(box[1] * scale) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
                cbx = [float(box[0] * scale), float(box[1] * scale), float(box[2] * scale), float(box[3] * scale), float(scores[i]), 0]
                car_box.append(cbx)
            submit_dicts.append({'image_id': image_id, 'car_bbox': car_box})
            # cv2.imwrite(os.path.join('pre', 'tfhead', image_id), src)
    f = open(export_resault, 'w')
    json.dump(submit_dicts, f, indent=1)
    f.close()

def get_cut_submit(images_path, model_weight, export_resault):
    classes = ['car']
    input_shape = [640, 640]
    num_classes = len(classes)
    print(num_classes)
    device = torch.device('cuda')
    net = yolov8.yolov8_detect(nc=num_classes, model_scale=0.33).to(device)
    net.load_state_dict(torch.load(model_weight, map_location=device))
    net.eval()

    images = os.listdir(images_path)
    submit_dicts = []
    pbar = tqdm(images)


    for image_id in pbar:
        pbar.set_description('get json | {}'.format(image_id))
        with torch.no_grad():
            image_path = os.path.join(images_path, image_id)
            src = cv2.imread(image_path)
            img = src.copy()
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            imgs, point = cut_image(img)
            point = torch.tensor(point)
            point = point.repeat((1, 2))
            imgs = [image_preprocess(x, (640, 640), True) for x in imgs]
            imgs = np.array(imgs)
            img_tensor = torch.from_numpy(imgs).float().to(device)
            img_tensor = img_tensor.permute(0, 3, 1, 2)

            y = net(img_tensor).detach().cpu()
            y = y.permute(0, 2, 1)             

            #y.shape batch, anchors, [cls xyxy]
            bboxes, cls_prefs = y.split((4, num_classes), 2)
            cls_confs, cls_pre = torch.max(cls_prefs, 2, keepdim=True)
            bboxes = torch.concat((bboxes, cls_confs, cls_pre), dim=2)
            nms_boxes = []
            for i, dbox in enumerate(bboxes):
                bbox, scores, _ = NMS(dbox, scores_threshold=0.4, device=torch.device('cpu'))
                nl = len(bbox)
                if nl == 0:
                    continue
                bbox = torch.cat(bbox, 0)
                bbox = bbox.reshape((-1, 4))

                p = point[i, :]
                p = p.repeat((len(bbox), 1))
                bbox = bbox + p
                scores = torch.tensor(scores).unsqueeze(-1)
                # scores = torch.cat(scores, 0)
                zeros = torch.zeros(scores.shape)
                box = torch.cat((bbox, scores, zeros), 1)
                nms_boxes.append(box)
            bbox = []
            scores = []
            car_box = []
            if len(nms_boxes) > 1:
                nms_boxes = torch.cat(nms_boxes, 0)
                bbox, scores, _ = min_NMS(nms_boxes, scores_threshold=0.5, device=torch.device('cpu'))
            elif len(nms_boxes) == 1:
                bbox = nms_boxes[0]
                box = bbox.squeeze(0)
                # print(nms_boxes)
                scores.append(box[4])
            for i, box in enumerate(bbox):
                src = cv2.rectangle(src, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0, 0.1))
                cbx = [float(box[0]), float(box[1]), float(box[2]), float(box[3]), float(scores[i]), 0]
                car_box.append(cbx)
            submit_dicts.append({'image_id': image_id, 'car_bbox': car_box})
            cv2.imwrite(os.path.join('pre', 'cut', image_id), src)
    f = open(export_resault, 'w')
    json.dump(submit_dicts, f, indent=1)
    f.close()

if __name__ == '__main__':
    # get_submit('../dataset/aisafety/new_test', 'weights/cbam/best_state_dict.pt', 'pre/jsons/new_cbam.json')
    get_submit('../dataset/aisafety/test', 'weights/tmp/epoch_60_val_loss=54.60627.pt', 'pre/jsons/n.json')
    # get_submit('../dataset/aisafety/new_test', 'weights/tfhead/best_state_dict.pt', 'pre/jsons/tfhead_0.5.json')
    # get_cut_submit('../dataset/aisafety/test', 'weights/cut_640/best_state_dict.pt', 'pre/jsons/cut_640.json')